import cv2
import numpy as np
from ultralytics import YOLO
import time
import cv2 as cv
import os
from datetime import datetime

# 摄像头矫正参数
k = [-0.281831359625869, 0.072276754058603, -0.007778520548340]
p = [-2.949785592530397e-04, -2.575583259892622e-04]
f = [9.796854694029854e+02, 9.797733387300718e+02]
c = [9.512391471508578e+02, 6.992317763795083e+02]

mtx = np.array([[f[0], 0.00000000e+00, c[0]],
                [0.00000000e+00, f[1], c[1]],
                [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]])
dist = np.array([k[0], k[1], p[0], p[1], k[2]])

# 全局变量
grid_size = 19  # 棋盘格数为 19x19
saved_points = None  # 保存第一次识别到的棋盘顶点
perspective_matrix = None  # 保存投影变换矩阵
last_valid_positions = {'black': [], 'white': []}  # 保存上次有效的棋子位置


def create_save_directory():
    """
    创建保存图片的目录
    """
    save_dir = "saved_chessboards"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    return save_dir


def save_warped_chessboard(warped_image, virtual_board=None):
    """
    保存投影变换后的棋盘图像
    """
    save_dir = create_save_directory()
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    # 保存投影变换后的棋盘
    warped_filename = f"warped_chessboard_{timestamp}.png"
    warped_path = os.path.join(save_dir, warped_filename)
    cv2.imwrite(warped_path, warped_image)
    print(f"✅ 投影变换后的棋盘已保存：{warped_path}")

    # 如果有虚拟棋盘，也一并保存
    if virtual_board is not None:
        virtual_filename = f"virtual_chessboard_{timestamp}.png"
        virtual_path = os.path.join(save_dir, virtual_filename)
        cv2.imwrite(virtual_path, virtual_board)
        print(f"✅ 虚拟棋盘已保存：{virtual_path}")

    return warped_path


def find_chessboard(image):
    """ 
    改进的棋盘识别函数，提高识别准确性
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # 使用高斯模糊减少噪声
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)

    # 使用自适应阈值处理
    adaptive_thresh = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                            cv2.THRESH_BINARY, 11, 2)

    # 使用多种边缘检测方法
    edges1 = cv2.Canny(blurred, 30, 100)
    edges2 = cv2.Canny(blurred, 50, 150)
    edges3 = cv2.Canny(blurred, 80, 200)

    # 组合边缘检测结果
    edges = cv2.bitwise_or(edges1, edges2)
    edges = cv2.bitwise_or(edges, edges3)

    # 形态学操作，闭合断开的边缘
    kernel = np.ones((3, 3), np.uint8)
    edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel, iterations=1)

    # 寻找轮廓
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # 候选棋盘列表
    candidates = []

    for contour in contours:
        area = cv2.contourArea(contour)
        perimeter = cv2.arcLength(contour, True)

        # 过滤掉太小的轮廓
        if area < 5000 or perimeter < 200:
            continue

        # 多边形逼近
        epsilon = 0.01 * perimeter  # 调整精度
        approx = cv2.approxPolyDP(contour, epsilon, True)

        # 如果是四边形
        if len(approx) == 4:
            # 检查是否为凸四边形
            if cv2.isContourConvex(approx):
                # 计算四边形的长宽比
                rect = cv2.minAreaRect(contour)
                width, height = rect[1]
                aspect_ratio = max(width, height) / min(width, height)

                # 棋盘应该接近正方形
                if aspect_ratio < 2.5:  # 长宽比不超过2.5
                    candidates.append({
                        'contour': approx,
                        'area': area,
                        'aspect_ratio': aspect_ratio,
                        'score': area / (aspect_ratio * 1000)  # 综合评分
                    })

    if not candidates:
        print("未检测到符合条件的棋盘轮廓！")
        return None

    # 按综合评分排序，选择最好的候选
    candidates.sort(key=lambda x: x['score'], reverse=True)

    # 进一步验证最佳候选
    best_candidate = candidates[0]
    best_contour = best_candidate['contour']

    # 验证四个角点的角度
    if validate_corner_angles(best_contour):
        print(f"找到棋盘！面积: {best_candidate['area']:.0f}, 长宽比: {best_candidate['aspect_ratio']:.2f}")
        return best_contour.reshape(4, 2)

    # 如果最佳候选不合格，尝试其他候选
    for candidate in candidates[1:]:
        if validate_corner_angles(candidate['contour']):
            print(f"找到棋盘！面积: {candidate['area']:.0f}, 长宽比: {candidate['aspect_ratio']:.2f}")
            return candidate['contour'].reshape(4, 2)

    print("未找到合适的棋盘！")
    return None


def validate_corner_angles(contour):
    """
    验证四边形的角度是否合理
    """
    points = contour.reshape(4, 2)

    # 计算四个角的角度
    angles = []
    for i in range(4):
        p1 = points[i]
        p2 = points[(i + 1) % 4]
        p3 = points[(i + 2) % 4]

        # 计算向量
        v1 = p1 - p2
        v2 = p3 - p2

        # 计算角度
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        cos_angle = np.clip(cos_angle, -1, 1)  # 防止数值误差
        angle = np.arccos(cos_angle) * 180 / np.pi
        angles.append(angle)

    # 检查角度是否合理（接近90度）
    for angle in angles:
        if angle < 45 or angle > 135:  # 角度应该在45-135度之间
            return False

    return True


def find_chessboard_with_lines(image):
    """
    基于直线检测的棋盘识别方法（备用）
    """
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)

    # 霍夫直线变换
    lines = cv2.HoughLines(edges, 1, np.pi / 180, threshold=100)

    if lines is None:
        return None

    # 分离水平线和垂直线
    horizontal_lines = []
    vertical_lines = []

    for line in lines:
        rho, theta = line[0]
        if abs(theta) < np.pi / 4 or abs(theta - np.pi) < np.pi / 4:
            horizontal_lines.append(line)
        elif abs(theta - np.pi / 2) < np.pi / 4:
            vertical_lines.append(line)

    # 如果检测到足够的水平线和垂直线，可能是棋盘
    if len(horizontal_lines) >= 3 and len(vertical_lines) >= 3:
        print(f"检测到 {len(horizontal_lines)} 条水平线和 {len(vertical_lines)} 条垂直线")
        # 这里可以进一步处理，找到棋盘边界
        # 暂时返回None，使用主要方法

    return None


def order_points(points):
    """
    对棋盘的四个顶点进行排序：左上、右上、右下、左下
    """
    rect = np.zeros((4, 2), dtype="float32")
    s = points.sum(axis=1)
    rect[0] = points[np.argmin(s)]
    rect[2] = points[np.argmax(s)]
    diff = np.diff(points, axis=1)
    rect[1] = points[np.argmin(diff)]
    rect[3] = points[np.argmax(diff)]
    return rect


def perform_perspective_transform(image, points=None):
    """
    对棋盘进行投影变换
    如果points为None，则使用全局保存的变换矩阵
    """
    global perspective_matrix

    if points is not None:
        # 第一次识别，计算变换矩阵
        ordered_points = order_points(points)
        width = 400
        height = 400
        dst_points = np.array([
            [0, 0],
            [width - 1, 0],
            [width - 1, height - 1],
            [0, height - 1]
        ], dtype="float32")
        perspective_matrix = cv2.getPerspectiveTransform(ordered_points, dst_points)

    # 使用保存的变换矩阵进行投影变换
    if perspective_matrix is not None:
        warped_image = cv2.warpPerspective(image, perspective_matrix, (400, 400))
        return warped_image
    else:
        return None


def generate_virtual_board(black_positions, white_positions):
    """
    根据棋盘信息生成虚拟棋盘
    """
    board_size = 400
    cell_size = board_size // (grid_size - 1)
    virtual_board = np.ones((board_size, board_size, 3), dtype=np.uint8) * 200

    for i in range(grid_size):
        cv2.line(virtual_board, (0, i * cell_size), (board_size, i * cell_size), (0, 0, 0), 1)
        cv2.line(virtual_board, (i * cell_size, 0), (i * cell_size, board_size), (0, 0, 0), 1)

    for x, y in black_positions:
        if 0 <= x < grid_size and 0 <= y < grid_size:
            cv2.circle(virtual_board, (x * cell_size, y * cell_size), 10, (0, 0, 0), -1)

    for x, y in white_positions:
        if 0 <= x < grid_size and 0 <= y < grid_size:
            cv2.circle(virtual_board, (x * cell_size, y * cell_size), 10, (255, 255, 255), -1)

    return virtual_board


def detect_pieces(image, model, matrix):
    """
    使用 YOLO 模型检测棋子，并计算棋子在棋盘上的落点
    """
    global last_valid_positions

    results = model(image)
    black_pieces = []
    white_pieces = []

    for result in results:
        for box in result.boxes:
            cls = int(box.cls)
            x_center, y_center, _, _ = box.xywh[0]
            x_center, y_center = int(x_center), int(y_center)

            if cls == 0:  # 黑子
                black_pieces.append((x_center, y_center))
            elif cls == 1:  # 白子
                white_pieces.append((x_center, y_center))

    # 输出所有黑白棋子的中心像素坐标
    print("黑子中心坐标：", black_pieces)
    print("白子中心坐标：", white_pieces)

    # 计算棋子在棋盘上的落点，并统计错误数量
    black_positions, black_error_count = calculate_positions(image, black_pieces, matrix, "黑子")
    white_positions, white_error_count = calculate_positions(image, white_pieces, matrix, "白子")

    total_error_count = black_error_count + white_error_count

    # 如果错误棋子数量超过3个，使用上次的有效位置
    if total_error_count >= 3:
        print(f"⚠️ 检测到{total_error_count}个错误棋子，超过阈值(3个)，使用上次有效数据")
        black_positions = last_valid_positions['black'].copy()
        white_positions = last_valid_positions['white'].copy()
        print("使用上次有效的黑子棋盘坐标：", black_positions)
        print("使用上次有效的白子棋盘坐标：", white_positions)
    else:
        # 更新上次有效位置
        last_valid_positions['black'] = black_positions.copy()
        last_valid_positions['white'] = white_positions.copy()
        print("✅ 棋子位置有效，已更新数据")
        print("黑子棋盘坐标：", black_positions)
        print("白子棋盘坐标：", white_positions)

    return black_positions, white_positions


def calculate_positions(image, pieces, matrix, piece_type):
    """
    计算棋子在棋盘上的落点，并返回错误数量
    """
    positions = []
    error_count = 0

    for x, y in pieces:
        # 将像素坐标转换为棋盘坐标
        board_x = round(x / (400 / (grid_size - 1)))
        board_y = round(y / (400 / (grid_size - 1)))

        # 确保坐标在棋盘范围内
        board_x = max(0, min(grid_size - 1, board_x))
        board_y = max(0, min(grid_size - 1, board_y))

        # 计算棋盘落点的像素坐标
        expected_x = int(board_x * (400 / (grid_size - 1)))
        expected_y = int(board_y * (400 / (grid_size - 1)))

        # 计算偏移距离
        distance = np.sqrt((x - expected_x) ** 2 + (y - expected_y) ** 2)

        # 偏移阈值（可以根据需要调整）
        threshold = 6  # 单位：像素

        if distance > threshold:
            error_count += 1
            # 在图像中标记偏离的棋子
            cv2.circle(image, (x, y), 10, (0, 255, 255), 2)
            cv2.putText(image, "Error", (x - 20, y - 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 255), 1)
            print(f"{piece_type}放落位置出错，请重新放置！ 棋子像素坐标: ({x}, {y}), 偏移距离: {distance:.2f}px")
        else:
            positions.append((board_x, board_y))

    return positions, error_count


def main():
    global saved_points, perspective_matrix, last_valid_positions

    model = YOLO("best.pt")
    # 打开视频流
    cap = cv.VideoCapture(0)
    cap.set(cv.CAP_PROP_FRAME_WIDTH, 2048)
    cap.set(cv.CAP_PROP_FRAME_HEIGHT, 1536)

    if not cap.isOpened():
        print("无法打开摄像头！")
        return

    last_update_time = time.time()
    chessboard_locked = False
    detection_attempts = 0  # 检测尝试次数
    max_attempts = 10  # 最大尝试次数

    # 存储当前帧的图像用于保存
    current_warped_image = None
    current_virtual_board = None

    print("控制键说明：")
    print("- 按 'q' 键退出程序")
    print("- 按 'r' 键重新识别棋盘")
    print("- 按 's' 键保存当前的投影变换后的棋盘")


    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法读取视频帧！")
            break

        # 每秒取一帧
        current_time = time.time()
        if current_time - last_update_time < 1:
            continue
        last_update_time = current_time

        # 摄像头矫正
        h, w = frame.shape[:2]
        newcameramtx, roi = cv.getOptimalNewCameraMatrix(
            mtx, dist, (w, h), 1, (w, h))
        corrected_frame = cv.undistort(frame, mtx, dist, None, mtx)

        # 裁剪图像
        display_size = (800, 600)
        dst_resized = cv.resize(corrected_frame, display_size, interpolation=cv.INTER_AREA)

        cv.imshow('dst', dst_resized)

        # 如果还没有锁定棋盘，尝试检测棋盘
        if not chessboard_locked:
            detection_attempts += 1
            print(f"正在尝试识别棋盘... (第 {detection_attempts} 次)")

            points = find_chessboard(dst_resized)

            if points is not None:
                # 第一次识别到棋盘，保存顶点坐标
                saved_points = points
                chessboard_locked = True
                print("棋盘已锁定！使用固定的四个顶点进行投影变换。")
                print(f"保存的棋盘顶点坐标：{saved_points}")

                # 在原图上标记检测到的棋盘顶点
                temp_image = dst_resized.copy()
                for i, point in enumerate(saved_points):
                    cv2.circle(temp_image, tuple(point.astype(int)), 5, (0, 255, 0), -1)
                    cv2.putText(temp_image, str(i + 1), tuple(point.astype(int) + [10, -10]),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
                cv2.imshow('Detected Chessboard', temp_image)
                cv2.waitKey(2000)  # 显示2秒
                cv2.destroyWindow('Detected Chessboard')

                # 计算并保存投影变换矩阵
                warped_image = perform_perspective_transform(dst_resized, saved_points)
                detection_attempts = 0
            else:
                if detection_attempts >= max_attempts:
                    print("多次尝试后仍未找到棋盘，请检查棋盘是否在视野内且光照充足")
                    detection_attempts = 0
                continue
        else:
            # 使用保存的顶点坐标进行投影变换
            warped_image = perform_perspective_transform(dst_resized)

        # 如果投影变换成功
        if warped_image is not None:
            # 棋子检测
            black_positions, white_positions = detect_pieces(warped_image, model, None)

            # 生成虚拟棋盘
            virtual_board = generate_virtual_board(black_positions, white_positions)

            # 更新当前帧的图像
            current_warped_image = warped_image.copy()
            current_virtual_board = virtual_board.copy()

            # 显示投影变换后的棋盘和虚拟棋盘
            cv.imshow("Warped Chessboard", warped_image)
            cv.imshow("Virtual Chessboard", virtual_board)

        # 按键处理
        key = cv.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('r'):
            # 按下 'r' 键重新识别棋盘
            chessboard_locked = False
            saved_points = None
            perspective_matrix = None
            detection_attempts = 0
            last_valid_positions = {'black': [], 'white': []}  # 重置有效位置
            print("重新识别棋盘模式已启用，请调整棋盘位置...")
        elif key == ord('s'):
            # 按下 's' 键保存当前的投影变换后的棋盘
            if current_warped_image is not None:
                saved_path = save_warped_chessboard(current_warped_image, current_virtual_board)
                print(f"📸 棋盘图像已保存！")
            else:
                print("❌ 没有可保存的棋盘图像！")

    cap.release()
    cv.destroyAllWindows()

]
if __name__ == "__main__":
    main()